import csv
import re


def extract_ratio(ratio_str):
    match = re.match(r'0?1:(\d+)', ratio_str)
    if match:
        return float(match.group(1))
    return None

data = []
with open("baoshan-schools-2024.csv", mode="r", encoding="utf-8") as file:
    reader = csv.reader(file)
    header = next(reader)
    for row in reader:
        if row and row[0] != "" and row[2] == "公办" and re.match(r'0?1:\d+', row[16]):
            data.append(row)

# 标准化列名
header_mapping = {
    "序号": "serial_no",
    "学校名称 / 学校": "school_name",
    "学校性质": "school_type",
    "教职工人数": "total_staff_count",
    "专职教师人数": "full_time_teacher_count",
    "师生比": "student_faculty_ratio_raw"
}
header = [header_mapping.get(col, col) for col in header]

clean_data = []
for row in data:
    serial_no = int(row[0])
    school_name = row[1]
    total_staff_count = int(row[14])
    full_time_teacher_count = int(row[15])
    student_faculty_ratio_raw = row[16]
    student_faculty_ratio = extract_ratio(student_faculty_ratio_raw)
    if student_faculty_ratio is None:
        continue

    estimated_student_count = int(round(total_staff_count * student_faculty_ratio))
    if full_time_teacher_count == 0:
        continue

    clean_row = [
        serial_no,
        school_name,
        total_staff_count,
        full_time_teacher_count,
        estimated_student_count
    ]
    clean_data.append(clean_row)

# 写入最终列名
final_header = [
    "serial_no",
    "school_name",
    "total_staff_count",
    "full_time_teacher_count",
    "estimated_student_count"
]

with open("clean-baoshan-schools-2024.csv", mode="w", encoding="utf-8", newline='') as file:
    writer = csv.writer(file)
    writer.writerow(final_header)
    writer.writerows(clean_data)